CREATE OR REPLACE FUNCTION cluster_reads(input_query text, output_table text) RETURNS VOID AS $$

# -----------------------------[ cluster_reads.sql ]---------------------------
"""
Groups overalapping reads, with the same orientation, into clusters.

USAGE:
SELECT cluster_reads(
  'SELECT chr, strand, read_start, read_end 
     FROM my_table 
     ORDER BY chr, strand, read_start', 
   'my_output_table'
   );

INPUT:
A table with columns: 
    1. chromosome, 
    2. strand orientation, 
    3. read start, 
    4. read end

Important: This table has to be sorted by chromosome, strand, read start.

Column names don't matter. Column order does. If different from above, 
change indexes as appropriate.
Columns other than the above are ignored.

OUTPUT:
A table like the input one with a column of tag clusters added.

"""

import psycopg2

#---------------------------------[ User's input ]-----------------------------

## column indexes. Change as appropriate
chr_column= 0
strand_column = 1
read_start_column= 2
read_end_column= 3

## Cluster reads overalpping by at least...
overlap_length = 0  ## +ve : Cluster reads overlapping by at least n bases
                    ## -ve : Cluster reads with a gap of no more than -n bases
                    ## 0 : Cluster reads where, at least, the end of one read
                    ##     correspond to the start of the next. I.e. no gap,
                    ##     no overlap. Adjacent reads.
                       
# ----------------------------[ Connect to database ]--------------------------

conn = psycopg2.connect("""dbname='CAGE' user='dberaldi' \
    host='localhost' password='mypassword'""")

cur_fetch = conn.cursor()
cur_modify= conn.cursor()

#-----------------------[ Output table ]---------------------------------------

cur_modify.execute("""CREATE TEMP TABLE tmp_cluster_table AS (
  %(input_query)s) """ %{'input_query':input_query} )

cur_modify.execute("""CREATE TABLE %(output_table)s AS (
  SELECT *, 'cluster' AS tag_cluster 
  FROM tmp_cluster_table
  WHERE 1=2) """ %{'output_table':output_table})

cur_modify.execute("""
 alter table %(output_table)s alter column tag_cluster TYPE varchar
  """ %{'output_table':output_table})

cur_fetch.execute(input_query)

# -----------------------------------------------------------------------------

line= cur_fetch.fetchone()
line= list(line)
line= [str(x) for x in line]

previous_chr= line[chr_column] + line[strand_column]
previous_end= line[read_end_column]

## Cluster name composed by: chromosome, strand, leftmost position of the cluster.
cluster_name= line[chr_column] + '_' + line[strand_column] + '_' + line[read_start_column]

line.append(cluster_name)

line_cluster= str(line).strip('[]')

cur_modify.execute("""
  INSERT INTO %(output_table)s 
  VALUES (%(line_cluster)s) 
  """ %{'output_table':output_table, 'line_cluster':line_cluster})

while 1<2:
    line= cur_fetch.fetchone()

    if line == None:
        break

    line= list(line)
    line= [str(x) for x in line]
    
    read_end= line[read_end_column]
    read_start= line[read_start_column]

    overlap= int(previous_end) - int(read_start)

    ## Update cluster_name when a gap between reads is found or when the read orientation changes    
    if  overlap >= (overlap_length - 1) and line[chr_column] + line[strand_column] == previous_chr:
        pass
    else:
        cluster_name = line[chr_column] + '_' + line[strand_column] + '_' + line[read_start_column]

    ## This 'if' makes small reads, contained into large ones, are clustered together
    if int(previous_end) < int(line[read_end_column]) or line[chr_column] != previous_chr:
        previous_end= line[read_end_column]
        
    line.append(cluster_name)

    line_cluster= str(line).strip('[]')

    cur_modify.execute("""
      INSERT INTO %(output_table)s 
      VALUES (%(line_cluster)s) 
      """ %{'output_table':output_table, 'line_cluster':line_cluster})

    previous_chr = line[chr_column] + line[strand_column]
      
#------------------------------------------------------------------------------

cur_modify.execute(""" DROP TABLE tmp_cluster_table """)

conn.commit()

$$ LANGUAGE 'plpythonu';